package com.boost.tjvm.compiler.ast;


import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.ProcessingEnvironment;
import javax.annotation.processing.RoundEnvironment;
import javax.annotation.processing.SupportedAnnotationTypes;
import javax.annotation.processing.SupportedSourceVersion;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.Element;
import javax.lang.model.element.ElementKind;
import javax.lang.model.element.TypeElement;
import javax.tools.Diagnostic;

import com.sun.source.util.Trees;
import com.sun.tools.javac.processing.JavacProcessingEnvironment;
import com.sun.tools.javac.tree.JCTree;
import com.sun.tools.javac.tree.TreeMaker;
import com.sun.tools.javac.tree.TreeTranslator;
import com.sun.tools.javac.tree.JCTree.JCAssert;
import com.sun.tools.javac.tree.JCTree.JCExpression;
import com.sun.tools.javac.tree.JCTree.JCStatement;
import com.sun.tools.javac.util.*;

import java.util.Set;


@SupportedAnnotationTypes("*")
@SupportedSourceVersion(SourceVersion.RELEASE_6)
public class ForceAssertions extends AbstractProcessor {
    private Trees trees;
    private TreeMaker make;
    private Name.Table names;

    @Override
    public synchronized void init(ProcessingEnvironment env) {
        super.init(env);
        trees = Trees.instance(env);
        Context context = ((JavacProcessingEnvironment) env).getContext();
        make = TreeMaker.instance(context);
        names = Names.instance(context).table;
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        if (!roundEnv.processingOver()) {
            Set<? extends Element> elements =
                    roundEnv.getRootElements();
            for (Element each : elements) {
                if (each.getKind() == ElementKind.CLASS) {
                    JCTree tree = (JCTree) trees.getTree(each);
                    TreeTranslator visitor = new Inliner();
                    tree.accept(visitor);
                }
            }
        } else
            processingEnv.getMessager().printMessage(
                    Diagnostic.Kind.NOTE,  " assertions  inlined.");
        return false;
    }

    private class Inliner extends TreeTranslator {

        @Override
        public void visitAssert(JCAssert tree) {
            super.visitAssert(tree);
            JCStatement newNode = makeIfThrowException(tree);
            result = newNode;
        }

        @Override
        public void visitMethodDef(JCTree.JCMethodDecl md) {
            super.visitMethodDef(md);



            String methodName = md.getName().toString();

            if(!methodName.equals("demo")){

                return;
            }

            JCTree.JCBlock block = md.getBody();
            List<JCTree.JCStatement> statements = block.stats;

            JCTree.JCExpression logMethod = chainDotsString(names, "info");
            JCTree.JCExpression logType = chainDotsString(names, "com/boost/tjvm/compiler/ast/Log");

            List<JCTree.JCVariableDecl> parameters = md.getParameters();

            JCTree.JCExpression apply = make.Apply(List.<JCTree.JCExpression>of(logType), logMethod,
                    generateLogArgs(parameters, "com/boost/tjvm/compiler/ast/ASSERT_CODE", methodName, make, md));

            ListBuffer<JCStatement> listBuffer = new ListBuffer<JCTree.JCStatement>();
            listBuffer.append(make.Exec(apply));

            for (JCTree.JCStatement stat : statements) {
                listBuffer.append(stat);
            }
            md.body.stats = listBuffer.toList();
            //annotationNode.getAst().setChanged();

            //JCStatement newNode = makeIfThrowException(tree);
            //result = newNode;
        }

        private JCStatement makeIfThrowException(JCAssert node) {
            // make: if (!(condition) throw new AssertionError(detail);
            List<JCExpression> args = node.getDetail() == null
                    ? List.<JCExpression>nil()
                    : List.of(node.detail);

            JCExpression expr = make.NewClass(
                    null,
                    null,
                    make.Ident(names.fromString("AssertionError")),
                    args,
                    null);
            return make.If(
                    make.Unary(JCTree.Tag.NOT, node.cond),
                    make.Throw(expr),
                    null);
        }

        /**
         * 生成log的参数表达式
         */
        public  List<JCTree.JCExpression> generateLogArgs(List<JCTree.JCVariableDecl> parameters, String className, String methodName, TreeMaker maker,JCTree.JCMethodDecl md) {
            JCTree.JCExpression[] argsArray = new JCTree.JCExpression[parameters.size() + 1];

            StringBuilder stringBuilder = new StringBuilder(className).append(".").append(methodName);
            if (parameters.size() > 0) {
                stringBuilder.append(" ");
                for (JCTree.JCVariableDecl variableDecl : parameters) {
                    stringBuilder.append(variableDecl.getName()).append(":{},");
                }
                stringBuilder.deleteCharAt(stringBuilder.length() - 1);
            } else {
                stringBuilder.append(" begin");
            }

            argsArray[0] = maker.Literal(stringBuilder.toString());
            JCTree.JCExpression jsonStringMethod =chainDotsString(names,"com.alibaba.fastjson.JSON.toJSONString");

            for (int i = 0; i < parameters.size(); i++) {
                argsArray[i + 1] = maker.Apply(List.<JCTree.JCExpression>nil(), jsonStringMethod, List.<JCTree.JCExpression>of(maker.Ident(parameters.get(i))));
            }
            return List.<JCTree.JCExpression>from(argsArray);
        }




        public JCExpression chainDotsString(Name.Table names, String elems) {
            return chainDots(names, null, null, elems.split("\\."));
        }

        public  JCExpression chainDots(Name.Table names, String elem1, String elem2, String... elems) {
            return chainDots(names, -1, elem1, elem2, elems);
        }

        public  JCExpression chainDots(Name.Table names, int pos, String elem1, String elem2, String... elems) {
            assert elems != null;
            TreeMaker maker = make;// node.getTreeMaker();
            //if (pos != -1) maker = make.at(pos);

            JCExpression e = null;
            if (elem1 != null) e = maker.Ident(names.fromString(elem1));
            if (elem2 != null) e = e == null ? maker.Ident(names.fromString(elem2)) : maker.Select(e, names.fromString(elem2));
            for (int i = 0 ; i < elems.length ; i++) {
                e = e == null ? maker.Ident(names.fromString(elems[i])) : maker.Select(e, names.fromString(elems[i]));
            }
            assert e != null;
            return e;
        }

    }

}